"""
Wave PDE implementation

The wave equation describes wave propagation:
    ∂²u/∂t² = c²∇²u

This second-order PDE is implemented as a first-order system:
    ∂u/∂t = v
    ∂v/∂t = c²∇²u
where u is the wave amplitude and v is the time derivative (velocity).
"""

using ..Fields
using ..Grids

"""
    WavePDE <: AbstractPDE

Wave equation for acoustic, electromagnetic, or mechanical wave propagation.

The mathematical definition as a second-order PDE is:
    ∂²u/∂t² = c²∇²u

This is implemented as a first-order system:
    ∂u/∂t = v
    ∂v/∂t = c²∇²u

where u is the wave field and v = ∂u/∂t is the velocity field.

# Parameters
- `speed`: Wave propagation speed c (default: 1.0)
- `bc`: Boundary conditions for wave field u (default: "auto_periodic_neumann")

# Properties
- Second-order hyperbolic PDE (as system: first-order)
- Conserves wave energy: E = ∫[v²/2 + c²|∇u|²/2] dV
- Supports wave propagation, reflection, and interference
- Characteristic speeds: ±c
"""
struct WavePDE <: AbstractPDE
    speed::Float64
    bc::String
    
    function WavePDE(; speed=1.0, bc="auto_periodic_neumann")
        new(speed, bc)
    end
end

"""
    evolution_rate(pde::WavePDE, state::FieldCollection, t::Float64)

Compute evolution rates for the wave equation system.

Expects state = [u, v] where:
- u: wave field (displacement)
- v: velocity field (∂u/∂t)

Returns [∂u/∂t, ∂v/∂t] = [v, c²∇²u]
"""
function evolution_rate(pde::WavePDE, state::FieldCollection, t::Float64)
    if length(state) != 2
        throw(ArgumentError("WavePDE requires exactly 2 fields: [u, v]"))
    end
    
    u = state[1]  # Wave field (displacement)
    v = state[2]  # Velocity field (∂u/∂t)
    
    # Time derivatives
    dudt = v  # ∂u/∂t = v
    
    # Wave equation: ∂v/∂t = c²∇²u
    u_laplacian = laplacian(u)
    dvdt = (pde.speed^2) * u_laplacian
    
    return FieldCollection([dudt, dvdt])
end

"""
    get_initial_condition(pde::WavePDE, u::ScalarField, v::Union{ScalarField, Nothing}=nothing)

Create initial condition for wave equation from displacement and velocity fields.

# Arguments
- `u`: Initial displacement field
- `v`: Initial velocity field (defaults to zero if not provided)

# Returns
- `FieldCollection([u, v])`: Combined initial state
"""
function get_initial_condition(pde::WavePDE, u::ScalarField, v::Union{ScalarField, Nothing}=nothing)
    if v === nothing
        v = ScalarField(zeros(size(data(u))), grid(u))
    end
    return FieldCollection([u, v], labels=["u", "v"])
end

"""
    wave_energy(pde::WavePDE, state::FieldCollection)

Compute the total wave energy: E = ∫[v²/2 + c²|∇u|²/2] dV

# Arguments
- `state`: FieldCollection containing [u, v]

# Returns
- Total wave energy (kinetic + potential)
"""
function wave_energy(pde::WavePDE, state::FieldCollection)
    u, v = state[1], state[2]
    
    # Kinetic energy: ∫v²/2 dV
    kinetic_energy = sum(data(v).^2) / 2
    
    # Potential energy: ∫c²|∇u|²/2 dV
    grad_u = gradient(u)
    potential_energy = 0.0
    
    if isa(grad_u, VectorField)
        grad_data = data(grad_u)
        if ndims(grad_data) >= 2
            for i in 1:size(grad_data, 1)
                potential_energy += sum(grad_data[i, :].^2)
            end
        else
            potential_energy += sum(grad_data.^2)
        end
    end
    potential_energy *= pde.speed^2 / 2
    
    return kinetic_energy + potential_energy
end

"""
    characteristic_speed(pde::WavePDE)

Get the characteristic wave speed for stability analysis.
"""
characteristic_speed(pde::WavePDE) = pde.speed

"""
    cfl_condition(pde::WavePDE, grid_spacing, dt)

Check Courant-Friedrichs-Lewy (CFL) condition for stability.
For wave equation: CFL = c·dt/dx ≤ 1
"""
function cfl_condition(pde::WavePDE, grid_spacing, dt)
    cfl = pde.speed * dt / grid_spacing
    return cfl, cfl <= 1.0
end

"""
    is_hyperbolic(pde::WavePDE)

Wave equation is hyperbolic.
"""
is_hyperbolic(pde::WavePDE) = true

"""
    requires_cfl_stability(pde::WavePDE)

Wave equations require CFL stability condition for explicit methods.
"""
requires_cfl_stability(pde::WavePDE) = true

"""
    expression(pde::WavePDE)

Return string representation of wave PDE system.
"""
function expression(pde::WavePDE)
    speed_str = pde.speed == 1.0 ? "" : "$(pde.speed)²"
    return "∂u/∂t = v, ∂v/∂t = $(speed_str)∇²u"
end

# Export wave PDE and analysis functions
export WavePDE, get_initial_condition, wave_energy,
       characteristic_speed, cfl_condition, is_hyperbolic, requires_cfl_stability